load("//xla:xla.default.bzl", "xla_cc_test")
load("//xla/tsl:tsl.bzl", "internal_visibility")
load(
    "//xla/tsl/platform:build_config.bzl",
    "tf_proto_library",
)
load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
load(
    "//xla/tsl/platform/default:cuda_build_defs.bzl",
    "if_cuda_is_configured",
)

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = internal_visibility(["//xla:internal"]),
    licenses = ["notice"],
)

package_group(
    name = "friends",
    includes = [
        "//xla:friends",
        "//xla:internal",
    ],
    packages = [
        "//tensorflow/core/tfrt/ifrt/...",
        "//third_party/australis/...",
        "//third_party/gxlang/...",
        "//third_party/openxla_pjrt_plugin/...",
        "//third_party/py/jax/...",
        "//third_party/torch_tpu/...",
    ],
)

cc_library(
    name = "worker_thread",
    srcs = ["worker_thread.cc"],
    hdrs = ["worker_thread.h"],
    deps = [
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:env",
    ],
)

cc_library(
    name = "event_pool",
    srcs = ["event_pool.cc"],
    hdrs = ["event_pool.h"],
    deps = [
        "//xla:types",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "semaphore",
    srcs = ["semaphore.cc"],
    hdrs = ["semaphore.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:types",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:logging",
    ],
)

xla_cc_test(
    name = "semaphore_test",
    srcs = ["semaphore_test.cc"],
    deps = [
        ":semaphore",
        "//xla/hlo/testlib:test",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:env",
    ],
)

cc_library(
    name = "abstract_tracked_device_buffer",
    srcs = ["abstract_tracked_device_buffer.cc"],
    hdrs = ["abstract_tracked_device_buffer.h"],
    visibility = [
        "//third_party/torch_tpu/pjrt:__pkg__",
        "//xla:internal",
    ],
    deps = [
        ":device_event",
        ":pjrt_client",
        ":pjrt_future",
        ":raw_buffer",
        "//xla:util",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "common_pjrt_client",
    srcs = [
        "common_pjrt_client.cc",
        "host_to_device_transfer_manager.cc",
    ],
    hdrs = [
        "common_pjrt_client.h",
        "host_to_device_transfer_manager.h",
    ],
    visibility = internal_visibility([
        ":friends",
        "//third_party/torch_tpu:__pkg__",
    ]),
    deps = [
        ":abstract_tracked_device_buffer",
        ":async_work_runner",
        ":device_event",
        ":host_callback",
        ":pjrt_client",
        ":pjrt_future",
        ":raw_buffer",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
        "@tsl//tsl/profiler/lib:connected_traceme",
        "@tsl//tsl/profiler/lib:context_types_hdrs",
        "@tsl//tsl/profiler/lib:scoped_memory_debug_annotation",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_cc_test(
    name = "tracked_device_buffer_test",
    srcs = ["tracked_device_buffer_test.cc"],
    deps = [
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_future",
        ":pjrt_stream_executor_client",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/client:client_library",
        "//xla/client:local_client",
        "//xla/hlo/testlib:test",
        "//xla/service:cpu_plugin",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "local_device_state",
    srcs = ["local_device_state.cc"],
    hdrs = ["local_device_state.h"],
    deps = [
        ":event_pool",
        ":pjrt_common",
        ":semaphore",
        ":worker_thread",
        "//xla:util",
        "//xla/client:local_client",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/protobuf:error_codes_proto_impl_cc",
        "//xla/tsl/util:env_var",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "pjrt_api",
    srcs = ["pjrt_api.cc"],
    hdrs = ["pjrt_api.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/tsl/platform:errors",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@tsl//tsl/platform",
    ],
)

xla_cc_test(
    name = "pjrt_api_test",
    srcs = ["pjrt_api_test.cc"],
    deps = [
        ":pjrt_api",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_wrapper_impl",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/protobuf:error_codes_proto_impl_cc",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:status_matchers",
    ],
)

cc_library(
    name = "pjrt_client",
    srcs = ["pjrt_client.cc"],
    hdrs = ["pjrt_client.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        ":pjrt_future",
        ":pjrt_layout",
        ":utils",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/distributed:key_value_store_interface",
        "//xla/service:computation_placer_hdr",
        "//xla/service:hlo_cost_analysis",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/protobuf:coordination_service_proto_cc",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@tsl//tsl/platform:errors",
    ],
)

cc_library(
    name = "pjrt_client_utils",
    srcs = ["pjrt_client_utils.cc"],
    hdrs = ["pjrt_client_utils.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":pjrt_client",
        "//xla:shape_util",
        "//xla:util",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "pjrt_client_test_common",
    testonly = 1,
    srcs = ["pjrt_client_test.cc"],
    hdrs = ["pjrt_client_test.h"],
    deps = [
        ":pjrt_client",
        ":pjrt_compiler",
        "//xla:cpu_function_runtime",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_builder",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/parser:hlo_parser",
        "//xla/hlo/testlib:test",
        "//xla/tests:literal_test_util",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:statusor",
    ],
    alwayslink = 1,
)

cc_library(
    name = "pjrt_executable",
    srcs = ["pjrt_executable.cc"],
    hdrs = ["pjrt_executable.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_common",
        ":pjrt_layout",
        "//xla:debug_options_flags",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/ffi:execution_context",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:executable_metadata_proto_cc",
        "//xla/pjrt/proto:execute_options_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:compiler",
        "//xla/service:computation_layout",
        "//xla/service:hlo_cost_analysis",
        "//xla/service:hlo_proto_cc",
        "//xla/service:hlo_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "pjrt_executable_test",
    srcs = ["pjrt_executable_test.cc"],
    deps = [
        ":pjrt_executable",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/service:computation_placer_hdr",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "pjrt_device_description",
    hdrs = ["pjrt_device_description.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_common",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "pjrt_compiler",
    srcs = ["pjrt_compiler.cc"],
    hdrs = ["pjrt_compiler.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_device_description",
        ":pjrt_executable",
        "//xla/hlo/builder:xla_computation",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/pjrt/proto:topology_description_proto_cc",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@tsl//tsl/platform:fingerprint",
    ],
)

xla_cc_test(
    name = "pjrt_compiler_test",
    srcs = ["pjrt_compiler_test.cc"],
    deps = [
        ":pjrt_client",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        "//xla:shape_util",
        "//xla/hlo/builder:xla_computation",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
    ],
)

cc_library(
    name = "pjrt_common",
    srcs = ["pjrt_common.cc"],
    hdrs = ["pjrt_common.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/pjrt/proto:pjrt_value_type_proto_cc",
        "//xla/tsl/lib/gtl:int_type",
    ],
)

xla_cc_test(
    name = "pjrt_common_test",
    srcs = ["pjrt_common_test.cc"],
    deps = [
        ":pjrt_common",
        "//xla/pjrt/proto:pjrt_value_type_proto_cc",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:status_matchers",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "utils",
    srcs = ["utils.cc"],
    hdrs = ["utils.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":layout_mode",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/service:computation_placer",
        "//xla/service:hlo_proto_cc",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@tsl//tsl/platform:errors",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:platform_port",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "pjrt_layout",
    hdrs = ["pjrt_layout.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        "//xla:shape_util",
        "//xla/hlo/parser:hlo_parser",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "layout_mode",
    srcs = ["layout_mode.cc"],
    hdrs = ["layout_mode.h"],
    visibility = ["//xla:friends"],
    deps = [
        "//xla:shape_util",
        "//xla/hlo/parser:hlo_parser",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
    ],
)

cc_library(
    name = "metrics",
    srcs = ["metrics.cc"],
    hdrs = ["metrics.h"],
    deps = ["//xla/tsl/lib/monitoring:counter"],
)

tf_proto_library(
    name = "stream_executor_executable_proto",
    srcs = ["stream_executor_executable.proto"],
    protodeps = [
        "//xla/pjrt/proto:compile_options_proto",
        "//xla:xla_proto",
    ],
    visibility = ["//visibility:public"],
)

cc_library(
    name = "stream_executor_executable",
    srcs = ["stream_executor_executable.cc"],
    hdrs = ["stream_executor_executable.h"],
    deps = [
        ":host_memory_spaces",
        ":pjrt_common",
        ":pjrt_executable",
        ":stream_executor_executable_proto_cc",
        "//xla:shape_util",
        "//xla:util",
        "//xla/client:local_client",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:compiler",
        "//xla/service:executable",
        "//xla/service:hlo_proto_cc",
        "//xla/service:hlo_proto_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "pjrt_stream_executor_client",
    srcs = [
        "pjrt_stream_executor_client.cc",
        "se_raw_buffer.cc",
        "tracked_device_buffer.cc",
    ],
    hdrs = [
        "pjrt_stream_executor_client.h",
        "se_raw_buffer.h",
        "tracked_device_buffer.h",
    ],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":abstract_tracked_device_buffer",
        ":common_pjrt_client",
        ":device_event",
        ":event_pool",
        ":host_callback",
        ":host_memory_spaces",
        ":layout_mode",
        ":local_device_state",
        ":metrics",
        ":mlir_to_hlo",
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_future",
        ":pjrt_stream_executor_device_description",
        ":raw_buffer",
        ":semaphore",
        ":stream_executor_executable",
        ":thread_pool_async_work_runner",
        ":transpose",
        ":utils",
        "//xla:executable_run_options",
        "//xla:literal",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/client:local_client",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/distributed:protocol_proto_cc",
        "//xla/pjrt/dump",
        "//xla/pjrt/profiling:device_time_measurement",
        "//xla/pjrt/profiling:profiling_context",
        "//xla/service:buffer_assignment",
        "//xla/service:compiler",
        "//xla/service:computation_layout",
        "//xla/service:computation_placer",
        "//xla/service:dump",
        "//xla/service:executable",
        "//xla/service:generic_transfer_manager",
        "//xla/service:hlo_cost_analysis",
        "//xla/service:hlo_proto_cc",
        "//xla/service:maybe_owning_device_memory",
        "//xla/service:shaped_buffer",
        "//xla/service:transfer_manager",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:device_memory_allocator",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:status",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@tsl//tsl/platform:casts",
        "@tsl//tsl/platform:fingerprint",
        "@tsl//tsl/platform:platform_port",
        "@tsl//tsl/profiler/lib:connected_traceme",
        "@tsl//tsl/profiler/lib:context_types_hdrs",
        "@tsl//tsl/profiler/lib:traceme",
    ],
    alwayslink = 1,
)

xla_cc_test(
    name = "pjrt_stream_executor_client_test",
    srcs = ["pjrt_stream_executor_client_test.cc"],
    deps = [
        ":local_device_state",
        ":pjrt_client",
        ":pjrt_executable",
        ":pjrt_future",
        ":pjrt_stream_executor_client",
        "//xla:literal",
        "//xla:literal_comparison",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/client:client_library",
        "//xla/client:local_client",
        "//xla/hlo/builder:xla_builder",
        "//xla/hlo/testlib:test",
        "//xla/service:cpu_plugin",
        "//xla/service:platform_util",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:path",
    ],
)

cc_library(
    name = "mlir_to_hlo",
    srcs = ["mlir_to_hlo.cc"],
    hdrs = ["mlir_to_hlo.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:util",
        "//xla/client:executable_build_options",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/translate:stablehlo",
        "//xla/mlir/utils:error_util",
        "//xla/mlir_hlo:hlo_dialect_registration",
        "//xla/mlir_hlo:mhlo_passes",
        "//xla/mlir_hlo:stablehlo_extension_passes",
        "//xla/service/spmd/shardy:constants",
        "//xla/service/spmd/shardy:utils",
        "//xla/service/spmd/shardy/sdy_round_trip:pipelines",
        "//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_export",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:ArithDialect",
        "@llvm-project//mlir:BytecodeWriter",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:FuncExtensions",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:MLProgramDialect",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:ReconcileUnrealizedCasts",
        "@llvm-project//mlir:ShapeDialect",
        "@llvm-project//mlir:Support",
        "@llvm-project//mlir:Transforms",
        "@shardy//shardy/dialect/sdy/ir:dialect",
        "@shardy//shardy/dialect/sdy/ir:register",
        "@stablehlo//:chlo_ops",
        "@stablehlo//:register",
        "@stablehlo//:stablehlo_ops",
        "@stablehlo//:stablehlo_passes",
        "@stablehlo//:stablehlo_portable_api",
        "@stablehlo//:stablehlo_serialization",
        "@stablehlo//:version",
    ],
)

xla_cc_test(
    name = "mlir_to_hlo_test",
    srcs = ["mlir_to_hlo_test.cc"],
    deps = [
        ":mlir_to_hlo",
        "//xla/hlo/testlib:test",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@stablehlo//:stablehlo_portable_api",
    ],
)

cc_library(
    name = "pjrt_future",
    srcs = ["pjrt_future.cc"],
    hdrs = ["pjrt_future.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:logging",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/base:no_destructor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/meta:type_traits",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@com_google_absl//absl/utility",
    ],
)

xla_cc_test(
    name = "pjrt_future_test",
    srcs = ["pjrt_future_test.cc"],
    deps = [
        ":pjrt_future",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
    ],
)

cc_library(
    name = "host_memory_spaces",
    srcs = ["host_memory_spaces.cc"],
    hdrs = ["host_memory_spaces.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_client",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:fingerprint",
    ],
)

# Transitional forwarding target. Use pjrt/plugin/xla_cpu:xla_cpu_pjrt_client instead.
cc_library(
    name = "tfrt_cpu_pjrt_client",
    hdrs = ["tfrt_cpu_pjrt_client.h"],
    visibility = internal_visibility([
        "//xla:friends",
    ]),
    deps = [
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
    ],
)

cc_library(
    name = "lru_cache",
    hdrs = ["lru_cache.h"],
    visibility = internal_visibility([
        "//xla:friends",
    ]),
    deps = [
        "@com_google_absl//absl/base:log_severity",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:hash_container_defaults",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@tsl//tsl/platform:logging",
    ],
)

xla_cc_test(
    name = "lru_cache_test",
    srcs = ["lru_cache_test.cc"],
    deps = [
        ":lru_cache",
        "//xla/hlo/testlib:test",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/random",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "transpose",
    srcs = [
        "transpose.cc",
        "transpose_kernels.h",
    ],
    hdrs = ["transpose.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":lru_cache",
        "//xla:compiler_macros",
        "//xla:ef57",
        "//xla:permutation_util",
        "//xla:util",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@com_google_absl//absl/types:variant",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_cc_test(
    name = "transpose_test",
    srcs = ["transpose_test.cc"],
    deps = [
        ":transpose",
        "//xla:array",
        "//xla:permutation_util",
        "//xla:shape_util",
        "//xla:util",
        "//xla/hlo/testlib:test",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/protobuf:error_codes_proto_impl_cc",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
        "@tsl//tsl/platform:env",
        "@tsl//tsl/platform:statusor",
        "@tsl//tsl/platform:test_benchmark",
        "@tsl//tsl/platform:test_main",
    ],
)

cc_library(
    name = "pjrt_c_api_client",
    srcs = ["pjrt_c_api_client.cc"],
    hdrs = ["pjrt_c_api_client.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":mlir_to_hlo",
        ":pjrt_api",
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        ":pjrt_future",
        ":pjrt_layout",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/ffi:execution_context",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
        "//xla/mlir_hlo:mhlo_passes",
        "//xla/pjrt/c:pjrt_c_api_ffi_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/pjrt/c:pjrt_c_api_layouts_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_memory_descriptions_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_profiler_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_stream_extension_hdrs",
        "//xla/pjrt/distributed:key_value_store_interface",
        "//xla/pjrt/extensions/cross_host_transfers:pjrt_c_api_cross_host_transfer_extension",
        "//xla/pjrt/extensions/executable_metadata:executable_metadata_extension",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:topology_description_proto_cc",
        "//xla/service:computation_placer_hdr",
        "//xla/service:hlo_cost_analysis",
        "//xla/service:hlo_proto_cc",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/protobuf:coordination_service_proto_cc",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Support",
        "@tsl//tsl/platform:casts",
        "@tsl//tsl/platform:fingerprint",
    ],
)

xla_cc_test(
    name = "pjrt_c_api_client_test",
    srcs = ["pjrt_c_api_client_test.cc"],
    tags = [
        "nomsan",
    ],
    deps = [
        ":mlir_to_hlo",
        ":pjrt_api",
        ":pjrt_c_api_client",
        ":pjrt_client",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        "//xla:cpu_function_runtime",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/hlo/builder:xla_builder",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/parser:hlo_parser",
        "//xla/pjrt/c:pjrt_c_api_cpu_internal",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/service:computation_placer_hdr",
        "//xla/tests:literal_test_util",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@stablehlo//:version",
    ],
)

cc_library(
    name = "tf_pjrt_client",
    srcs = ["tf_pjrt_client.cc"],
    hdrs = ["tf_pjrt_client.h"],
    visibility = ["//visibility:public"],
    deps = [
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/service:computation_placer_hdr",
        "//xla/service:hlo_cost_analysis",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@tsl//tsl/platform:casts",
    ],
)

xla_cc_test(
    name = "tf_pjrt_client_test",
    srcs = ["tf_pjrt_client_test.cc"],
    deps = [
        ":tf_pjrt_client",
        "//xla:literal_util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/parser:hlo_parser",
        "//xla/pjrt/plugin/xla_cpu:cpu_client_options",
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
        "//xla/service:hlo_proto_cc",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:env",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "host_callback",
    srcs = ["host_callback.cc"],
    hdrs = ["host_callback.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_client",
        ":pjrt_executable",
        ":pjrt_future",
        "//xla:shape_util",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:ffi",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:status",
    ],
)

xla_cc_test(
    name = "host_callback_test",
    srcs = ["host_callback_test.cc"],
    deps = [
        ":host_callback",
        ":pjrt_client",
        "//xla:xla_data_proto_cc",
        "//xla/tests:literal_test_util",
        "//xla/tsl/lib/core:status_test_util",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "exceptions",
    hdrs = ["exceptions.h"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
    ],
)

cc_library(
    name = "status_casters",
    hdrs = ["status_casters.h"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":exceptions",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@tsl//tsl/platform:macros",
    ],
)

cc_library(
    name = "triton",
    srcs = if_cuda_is_configured(
        ["triton_cuda.cc"],
        ["triton_stub.cc"],
    ),
    hdrs = ["triton.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:logging",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@llvm-project//llvm:TargetParser",
        "@tsl//tsl/platform:cuda_root_path",
        "@tsl//tsl/platform:path",
    ] + if_cuda_is_configured(
        [
            "@com_google_absl//absl/base",
            "@com_google_absl//absl/status",
            "@com_google_absl//absl/strings:str_format",
            "@com_google_absl//absl/strings",
            "@llvm-project//llvm:Core",
            "@llvm-project//llvm:IRReader",
            "@llvm-project//llvm:Linker",
            "@llvm-project//llvm:MC",
            "@llvm-project//llvm:Support",
            "@llvm-project//llvm:Target",
            "@llvm-project//mlir:ArithDialect",
            "@llvm-project//mlir:ArithToLLVM",
            "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
            "@llvm-project//mlir:BytecodeWriter",
            "@llvm-project//mlir:ConversionPasses",
            "@llvm-project//mlir:DataLayoutInterfaces",
            "@llvm-project//mlir:ExecutionEngineUtils",
            "@llvm-project//mlir:FuncExtensions",
            "@llvm-project//mlir:IndexToLLVM",
            "@llvm-project//mlir:IR",
            "@llvm-project//mlir:LLVMCommonConversion",
            "@llvm-project//mlir:LLVMIRTransforms",
            "@llvm-project//mlir:LLVMToLLVMIRTranslation",
            "@llvm-project//mlir:NVVMTarget",
            "@llvm-project//mlir:NVVMToLLVMIRTranslation",
            "@llvm-project//mlir:Parser",
            "@llvm-project//mlir:Pass",
            "@llvm-project//mlir:SCFToControlFlow",
            "@llvm-project//mlir:ToLLVMIRTranslation",
            "@llvm-project//mlir:Transforms",
            "//xla/backends/gpu/codegen/triton:compilation_pipeline_cuda",
            "//xla/tsl/platform:errors",
            "//xla/tsl/platform:statusor",
            "@triton//:TritonDialects",
            "@triton//:TritonGPUToLLVM",
            "@triton//:TritonGPUTransforms",
            "@triton//:TritonNvidiaGPUTransforms",
            "@triton//:TritonToTritonGPU",
            "@triton//:TritonTransforms",
            "@triton//third_party/nvidia:NVGPUToLLVM",
            "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM",
        ],
        ["@com_google_absl//absl/status"],
    ),
)

cc_library(
    name = "device_event",
    hdrs = ["device_event.h"],
    deps = [
        ":pjrt_future",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "raw_buffer",
    srcs = ["raw_buffer.cc"],
    hdrs = ["raw_buffer.h"],
    deps = [
        ":async_work_runner",
        ":device_event",
        ":pjrt_client",
        ":pjrt_future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "raw_buffer_test",
    srcs = ["raw_buffer_test.cc"],
    deps = [
        ":raw_buffer",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "pjrt_stream_executor_device_description",
    hdrs = ["pjrt_stream_executor_device_description.h"],
    deps = [
        ":pjrt_device_description",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "async_work_runner",
    hdrs = ["async_work_runner.h"],
    deps = [
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "thread_pool_async_work_runner",
    srcs = ["thread_pool_async_work_runner.cc"],
    hdrs = ["thread_pool_async_work_runner.h"],
    deps = [
        ":async_work_runner",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "pjrt_phase_compile",
    srcs = ["pjrt_phase_compile.cc"],
    hdrs = ["pjrt_phase_compile.h"],
    visibility = ["//visibility:public"],
    deps = [
        ":pjrt_c_api_client",
        ":pjrt_compiler",
        ":pjrt_executable",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "pjrt_phase_compile_sample_plugin",
    srcs = ["pjrt_phase_compile_sample_plugin.cc"],
    hdrs = ["pjrt_phase_compile_sample_plugin.h"],
    visibility = ["//visibility:private"],
    deps = [
        ":pjrt_client",
        ":pjrt_compiler",
        ":pjrt_executable",
        "//xla/hlo/builder:xla_computation",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/c:pjrt_c_api_wrapper_impl",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Transforms",
        "@stablehlo//:stablehlo_passes_optimization",
        "@stablehlo//:stablehlo_portable_api",
        "@stablehlo//:stablehlo_serialization",
    ],
)

xla_cc_test(
    name = "pjrt_phase_compiler_test",
    srcs = ["pjrt_phase_compiler_test.cc"],
    deps = [
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_phase_compile_sample_plugin",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/testlib:test",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/plugin/xla_cpu:cpu_topology",
        "//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:status_matchers",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@stablehlo//:reference_api",
    ],
)

xla_cc_test(
    name = "pjrt_phase_compile_extension_test",
    srcs = ["pjrt_phase_compile_extension_test.cc"],
    deps = [
        ":pjrt_api",
        ":pjrt_c_api_client",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_phase_compile",
        ":pjrt_phase_compile_sample_plugin",
        "//xla/pjrt/c:pjrt_c_api_cpu_internal",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:status_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@stablehlo//:reference_api",
    ],
)
